package cn.zhxu.research;

import cn.zhxu.data.Mapper;
import cn.zhxu.okhttps.HTTP;
import cn.zhxu.okhttps.HttpResult;
import cn.zhxu.okhttps.OkHttps;
import cn.zhxu.toys.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.stereotype.Component;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.time.Duration;
import java.util.*;
import java.util.stream.Collectors;

import static cn.zhxu.research.DataItem.SCORE_MAX;

@Component
public class Translator implements ApplicationRunner {

    private static final Logger log = LoggerFactory.getLogger(Translator.class);

    private static final HTTP http = OkHttps.getHttp()
            .newBuilder()
            .config(b -> {
                b.connectTimeout(Duration.ofSeconds(60));
                b.writeTimeout(Duration.ofSeconds(60));
                b.readTimeout(Duration.ofSeconds(300));
            })
            .build();

    public static String translate(String text) {
        String content = "请将下面的每行文字都分别翻译为英文（只输出译文，不要输出其它内容）。\n\n" + text;
        Map<String, String> message = new HashMap<>();
        message.put("role", "user");
        message.put("content", content);
        var res = http.sync("https://open.bigmodel.cn/api/paas/v4/chat/completions")
                .bodyType(OkHttps.JSON)
                .bearerAuth("1c2ce75fe41b4ab1b1e88f55e9a3e05c.IexfHhRhGZ0vDVOQ")
                .addBodyPara("model", "glm-4-plus")
                .addBodyPara("messages", Collections.singletonList(message))
                .post();
        if (res.getState() != HttpResult.State.RESPONSED) {
            log.info(res.toString());
            throw new RuntimeException("Translate No Response!");
        }
        Mapper body = res.getBody().toMapper();
        int status = res.getStatus();
        if (status != 200) {
            if (status == 400) {
                String code = body.getMapper("error").getString("code");
                if ("1301".equals(code)) {
                    log.info("触发了 1301 内容限制，跳过该段数据:");
                    return null;
                }
                if ("1302".equals(code)) {
                    log.info("触发了 1302 并发限制，等待 10 秒后重试：{}", body);
                    try {
                        Thread.sleep(10000);
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                    return translate(text);
                }
            }
            log.info(res.toString());
            log.info(body.toString());
            throw new RuntimeException("Translate Response Status error: " + status);
        }
        try {
            return body.getArray("choices")
                    .getMapper(0)
                    .getMapper("message")
                    .getString("content");
        } catch (Exception e) {
            log.info(body.toString());
            throw new RuntimeException("Translate Response Body error");
        }
    }

    public static void main(String[] args) throws Exception {
        new Translator().run(null);
    }

    @Override
    public void run(ApplicationArguments args) throws Exception {
        log.info("开始翻译...");
        var dataItems = DataItem.readDataItems("E:/SpaceMine/research/src/main/resources/dataset_raw.txt");
        var totalCnt = dataItems.size();
        int processedCnt = 0;
        try (BufferedWriter skippedWriter = new BufferedWriter(new FileWriter("E:/SpaceMine/research/src/main/resources/dataset.skip.txt"))) {
            try (BufferedWriter writer = new BufferedWriter(new FileWriter("E:/SpaceMine/research/src/main/resources/dataset2.txt"))) {
                while (processedCnt < totalCnt) {
                    int count = Math.min(1, totalCnt - processedCnt);
                    log.info("Translate {} ~ {}", processedCnt, processedCnt + count);
                    var items = dataItems.subList(processedCnt, processedCnt + count);
                    var text = items.stream()
                            .map(it -> it.getTxt1() + "\n" + it.getTxt2())
                            .collect(Collectors.joining("\n"));
                    var results = preprocess(translate(text), count * 2);
                    if (results == null) {
                        writeToSkipped(skippedWriter, items);
                        processedCnt += count;
                        continue;
                    }
                    for (int i = 0; i < count; i++) {
                        var item = items.get(i);
                        var en1 = results[2 * i];
                        var en2 = results[2 * i + 1];
                        writer.write(item.getTxt1() + "##" + en1 + "##" + SCORE_MAX);
                        writer.newLine();
                        writer.write(item.getTxt1() + "##" + en2 + "##" + item.getScore());
                        writer.newLine();
                        writer.write(item.getTxt2() + "##" + en2 + "##" + SCORE_MAX);
                        writer.newLine();
                        writer.write(item.getTxt2() + "##" + en1 + "##" + item.getScore());
                        writer.newLine();
                    }
                    processedCnt += count;
                }
            }
        }
        log.info("翻译结束！");
    }

    static String[] preprocess(String translatedText, int expectedCount) {
        if (translatedText == null) {
            return null;
        }
        var texts = Arrays.stream(translatedText.split("\n"))
                .filter(StringUtils::notBlank)
                .toArray(String[]::new);
        if (texts.length != expectedCount) {
            log.info("翻译的数据条数不对，跳过该段数据：{} != {}", texts.length, expectedCount);
            return null;
        }
        return texts;
    }

    static void writeToSkipped(BufferedWriter writer, List<DataItem> items) throws IOException {
        for (var item : items) {
            writer.write(item.getTxt1() + "##" + item.getTxt2() + "##" + item.getScore());
            writer.newLine();
        }
    }

}
